from django.shortcuts import render, get_object_or_404
from rest_framework import viewsets, filters, status
from rest_framework.decorators import action
from rest_framework.response import Response
from django_filters.rest_framework import DjangoFilterBackend
from django.db.models import Q
from datetime import datetime, timedelta
from django.db import transaction

from .models import Stock, StockPrice
from .serializers import StockSerializer, StockPriceSerializer, StockDetailSerializer

class StockViewSet(viewsets.ModelViewSet):
    """股票信息视图集"""
    queryset = Stock.objects.all()
    serializer_class = StockSerializer
    filter_backends = [DjangoFilterBackend, filters.SearchFilter, filters.OrderingFilter]
    filterset_fields = ['industry', 'market', 'is_active']
    search_fields = ['code', 'name']
    ordering_fields = ['code', 'name', 'created_at', 'updated_at']
    
    def get_serializer_class(self):
        if self.action == 'list' or self.action == 'retrieve':
            return StockDetailSerializer
        return StockSerializer
        
    @action(detail=False, methods=['post'])
    def batch_create(self, request):
        """批量创建股票信息"""
        data = request.data
        if not isinstance(data, list):
            return Response({"error": "数据格式错误，应为股票信息列表"}, status=status.HTTP_400_BAD_REQUEST)
            
        serializer = self.get_serializer(data=data, many=True)
        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data, status=status.HTTP_201_CREATED)
        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
        
    @action(detail=True, methods=['patch'])
    def update_status(self, request, pk=None):
        """更新股票活跃状态"""
        stock = self.get_object()
        is_active = request.data.get('is_active')
        
        if is_active is None:
            return Response({"error": "缺少is_active参数"}, status=status.HTTP_400_BAD_REQUEST)
            
        stock.is_active = is_active
        stock.save()
        
        serializer = self.get_serializer(stock)
        return Response(serializer.data)
    
    @action(detail=True, methods=['get'])
    def prices(self, request, pk=None):
        """获取指定股票的价格历史数据"""
        stock = self.get_object()
        
        # 获取查询参数
        days = request.query_params.get('days')
        start_date = request.query_params.get('start_date')
        end_date = request.query_params.get('end_date')
        
        # 构建查询条件
        queryset = stock.prices.all()
        
        if days:
            try:
                days = int(days)
                end_date = datetime.now().date()
                start_date = end_date - timedelta(days=days)
                queryset = queryset.filter(date__gte=start_date, date__lte=end_date)
            except ValueError:
                pass
        elif start_date and end_date:
            try:
                start_date = datetime.strptime(start_date, '%Y-%m-%d').date()
                end_date = datetime.strptime(end_date, '%Y-%m-%d').date()
                queryset = queryset.filter(date__gte=start_date, date__lte=end_date)
            except ValueError:
                pass
        
        # 排序并序列化
        queryset = queryset.order_by('date')
        serializer = StockPriceSerializer(queryset, many=True)
        return Response(serializer.data)

class StockPriceViewSet(viewsets.ModelViewSet):
    """股票价格历史数据视图集"""
    queryset = StockPrice.objects.all()
    serializer_class = StockPriceSerializer
    filter_backends = [DjangoFilterBackend, filters.OrderingFilter]
    filterset_fields = ['stock', 'date']
    ordering_fields = ['date', 'close_price', 'volume']
    
    @action(detail=False, methods=['get'])
    def latest(self, request):
        """获取最新的股票价格数据"""
        # 获取查询参数
        stock_code = request.query_params.get('code')
        days = request.query_params.get('days', 1)
        
        try:
            days = int(days)
        except ValueError:
            days = 1
            
        # 构建查询条件
        queryset = self.get_queryset()
        
        if stock_code:
            queryset = queryset.filter(stock__code=stock_code)
            
        # 获取每只股票最新的价格数据
        latest_prices = []
        stocks = Stock.objects.filter(is_active=True)
        
        if stock_code:
            stocks = stocks.filter(code=stock_code)
            
        for stock in stocks:
            prices = queryset.filter(stock=stock).order_by('-date')[:days]
            serializer = StockPriceSerializer(prices, many=True)
            latest_prices.extend(serializer.data)
            
        return Response(latest_prices)
    
    @action(detail=False, methods=['post'])
    def batch_create(self, request):
        """批量创建股票价格数据"""
        data = request.data
        if not isinstance(data, list):
            return Response({"error": "数据格式错误，应为股票价格数据列表"}, status=status.HTTP_400_BAD_REQUEST)
        
        # 使用事务确保数据一致性
        with transaction.atomic():
            # 验证每条数据中的股票代码是否存在
            for item in data:
                stock_code = item.get('stock_code')
                if not stock_code:
                    return Response({"error": "每条数据必须包含stock_code字段"}, status=status.HTTP_400_BAD_REQUEST)
                
                try:
                    stock = Stock.objects.get(code=stock_code)
                    # 将stock_code替换为stock对象的id
                    item['stock'] = stock.id
                except Stock.DoesNotExist:
                    return Response({"error": f"股票代码 {stock_code} 不存在"}, status=status.HTTP_404_NOT_FOUND)
            
            serializer = self.get_serializer(data=data, many=True)
            if serializer.is_valid():
                serializer.save()
                return Response(serializer.data, status=status.HTTP_201_CREATED)
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
    
    @action(detail=False, methods=['post'])
    def update_or_create(self, request):
        """更新或创建股票价格数据"""
        data = request.data
        stock_code = data.get('stock_code')
        date = data.get('date')
        
        if not stock_code or not date:
            return Response({"error": "缺少必要参数stock_code或date"}, status=status.HTTP_400_BAD_REQUEST)
            
        try:
            stock = Stock.objects.get(code=stock_code)
        except Stock.DoesNotExist:
            return Response({"error": f"股票代码 {stock_code} 不存在"}, status=status.HTTP_404_NOT_FOUND)
            
        try:
            date_obj = datetime.strptime(date, '%Y-%m-%d').date()
        except ValueError:
            return Response({"error": "date格式错误，应为YYYY-MM-DD"}, status=status.HTTP_400_BAD_REQUEST)
            
        # 替换stock_code为stock对象的id
        data['stock'] = stock.id
        
        # 尝试查找现有记录
        try:
            stock_price = StockPrice.objects.get(stock=stock, date=date_obj)
            # 更新现有记录
            serializer = self.get_serializer(stock_price, data=data, partial=True)
        except StockPrice.DoesNotExist:
            # 创建新记录
            serializer = self.get_serializer(data=data)
            
        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data)
        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
        
    @action(detail=False, methods=['post'])
    def import_from_csv(self, request):
        """从CSV文件导入股票价格数据"""
        file = request.FILES.get('file')
        if not file:
            return Response({"error": "未提供CSV文件"}, status=status.HTTP_400_BAD_REQUEST)
        
        if not file.name.endswith('.csv'):
            return Response({"error": "文件格式错误，应为CSV文件"}, status=status.HTTP_400_BAD_REQUEST)
        
        # 使用事务确保数据一致性
        with transaction.atomic():
            try:
                # 解析CSV文件
                import csv
                from io import StringIO
                
                # 读取文件内容
                content = file.read().decode('utf-8')
                csv_file = StringIO(content)
                reader = csv.DictReader(csv_file)
                
                # 准备批量创建的数据
                stock_prices = []
                for row in reader:
                    stock_code = row.get('stock_code')
                    if not stock_code:
                        return Response({"error": "CSV文件中缺少stock_code列"}, status=status.HTTP_400_BAD_REQUEST)
                    
                    try:
                        stock = Stock.objects.get(code=stock_code)
                    except Stock.DoesNotExist:
                        return Response({"error": f"股票代码 {stock_code} 不存在"}, status=status.HTTP_404_NOT_FOUND)
                    
                    # 构建数据对象
                    stock_price_data = {
                        'stock': stock.id,
                        'date': row.get('date'),
                        'open_price': row.get('open_price'),
                        'close_price': row.get('close_price'),
                        'high_price': row.get('high_price'),
                        'low_price': row.get('low_price'),
                        'volume': row.get('volume')
                    }
                    
                    stock_prices.append(stock_price_data)
                
                # 批量创建数据
                serializer = self.get_serializer(data=stock_prices, many=True)
                if serializer.is_valid():
                    serializer.save()
                    return Response({"message": f"成功导入 {len(stock_prices)} 条股票价格数据"}, status=status.HTTP_201_CREATED)
                return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
                
            except Exception as e:
                return Response({"error": f"导入CSV文件失败: {str(e)}"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)